package ai.george.quickstart.advisor;


import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.Prompt;

import java.util.*;

/**
 * 自定义的Advisor
 * 实现了会话记忆的功能
 */
@Slf4j
public class CustomMessageChatMemoryAdvisor implements BaseAdvisor {

    private static final String SESSION_ID = "sessionId";

    /**
     * 会话记忆(这里只是用内存作为会话储存的一种演示，实际开发历史会话需要保存到数据库)
     * 注意：这里变量必须声明为 static 形式
     */
    private static Map<String, List<Message>>  messageMap = new HashMap<>();


    @Override
    public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {

        String sessionId = (String) chatClientRequest.context().get(SESSION_ID);
        List<Message> historyMessages = messageMap.get(sessionId);
        boolean isEmpty = historyMessages == null;
        if (isEmpty) {
            historyMessages = new ArrayList<>();
        }
        //获取本次请求的Message列表
        List<Message> thisRequestMessage = chatClientRequest.prompt().getInstructions();
        //将本次消息添加到历史消息中
        historyMessages.addAll(thisRequestMessage);
        if (isEmpty) {
            //只有第一次请求，才需要put,因为Map已经引用了这个List对象，addAll添加时，已经添加到列表中
            messageMap.put(sessionId, historyMessages);
        }

        //更新prompt对象(生成一份副本，并带上了历史所有消息)
        Prompt newPrompt = chatClientRequest.prompt().mutate()
                .messages(historyMessages)
                .build();
        //更新request对象
        ChatClientRequest newChatClientRequest = chatClientRequest.mutate()
                .prompt(newPrompt)
                .build();
        return newChatClientRequest;
    }

    @Override
    public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
        String sessionId = (String) chatClientResponse.context().get(SESSION_ID);
        List<Message> historyMessages = messageMap.get(sessionId);

        if (Objects.isNull(historyMessages)) {
            return chatClientResponse;
        }

        //获取本次响应的Message
        Message responseMessage = chatClientResponse.chatResponse().getResult().getOutput();
        historyMessages.add(responseMessage);
        return chatClientResponse;
    }

    @Override
    public int getOrder() {
        return 3;
    }
}
